# -*- coding: utf-8 -*-            
# @Author : HJH
# @Time : 2024/3/12 14:52
# @Use : 往返航线数据匹配
import hashlib
import itertools
import json
from datetime import datetime

import pandas as pd

if __name__ == '__main__':
    file_name = input("输入需要匹配文件名(.xlsx格式): ")
    out_file_name = input("输入需要匹配文件名(.xlsx格式): ")
    print("读取文件数据... ...")
    pd_data = pd.read_excel(f'./{file_name}.xlsx')
    # 将每一行转为字典，并存储在一个列表中
    data_dict_list = pd_data.to_dict(orient='records')


    def generate_group_key(item):
        # 定义一个函数来生成组合键
        # 假设你想根据 'field1' 和 'field2' 两个字段进行分组
        key_str = item['出发地'] + item['目的地'] + item['舱位'] + str(item['价格'])
        hash_object = hashlib.md5()
        hash_object.update(key_str.encode())
        return hash_object.hexdigest()


    # 假设 data_dict_list 是一个字典列表，包含了多个字段
    data_dict_list_sorted = sorted(data_dict_list, key=lambda x: x['出发地'])

    grouped_data = {}
    for item in data_dict_list:
        key = generate_group_key(item)
        if key not in grouped_data:
            grouped_data[key] = []
        grouped_data[key].append(item)

    data_dict_list = []
    for items in grouped_data.values():
        if len(items) == 1:
            items[0]['开始日期'] = items[0]['起飞日期']
            items[0]['结束日期'] = items[0]['起飞日期']
            data_dict_list.append(items[0])
        else:
            # 根据日期排序
            items_sorted = sorted(items, key=lambda x: x['起飞日期'])
            merged_lines = []
            for item in items_sorted:
                if len(merged_lines) > 0:

                    temp_date = datetime.strptime(merged_lines[-1]['起飞日期'], '%Y-%m-%d') + pd.Timedelta(days=1)
                    if temp_date.strftime('%Y-%m-%d') == item['起飞日期']:
                        if item == items_sorted[-1]:
                            first_item = merged_lines[0]

                            first_item['开始日期'] = first_item['起飞日期']
                            first_item['结束日期'] = item['起飞日期']
                            data_dict_list.append(first_item)
                        else:
                            merged_lines.append(item)
                    else:
                        first_item = merged_lines[0]
                        last_item = merged_lines[-1]

                        first_item['开始日期'] = first_item['起飞日期']
                        first_item['结束日期'] = last_item['起飞日期']
                        data_dict_list.append(first_item)

                        if item == items_sorted[-1]:
                            item['开始日期'] = item['起飞日期']
                            item['结束日期'] = item['起飞日期']
                            data_dict_list.append(item)
                        else:
                            merged_lines = [item]
                else:
                    merged_lines.append(item)

    res_data = []
    print("开始匹配处理数据... ...")
    for data_dict in data_dict_list:
        filter_data = [d for d in data_dict_list if
                       d['出发地'] == data_dict['目的地'] and d['目的地'] == data_dict['出发地']]
        for filter_item in filter_data:
            if data_dict['开始日期'] >= filter_item['开始日期']:
                new_item = dict(filter_item)
                new_item['B_出发地'] = data_dict['出发地']
                new_item['B_目的地'] = data_dict['目的地']
                new_item['B_开始日期'] = data_dict['开始日期']
                new_item['B_结束日期'] = data_dict['结束日期']
                new_item['B_舱位'] = data_dict['舱位']
                new_item['B_价格'] = data_dict['价格']
                new_item['B_加价后'] = int(data_dict['价格']) + 500
                res_data.append(new_item)

    # 将列表转换为 DataFrame
    df = pd.DataFrame(res_data)

    # # 设置表头
    # header = ['航班号', '出发日期', '']

    print("开始写入文件... ...")
    # 将数据写入 Excel 文件并设置表头
    with pd.ExcelWriter(f'./{out_file_name}.xlsx', engine='xlsxwriter') as writer:
        df.to_excel(writer, index=False, )

    # 输出成功提示
    print("Excel 文件写入成功！")
